#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
EndoSight-UC AI推理服务器 - 修复版本
修复模型推理问题，确保基于真实图像内容进行推理
"""

import os
import sys
import logging
import traceback
import time
from pathlib import Path
from typing import Dict, Any, Optional
import uuid
import json

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from torchvision.models import ResNet50_Weights
from transformers import BertModel, BertTokenizer
from PIL import Image
import numpy as np
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn
from pydantic import BaseModel

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('ai_inference_fixed.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

# 全局变量
model = None
device = None
tokenizer = None

# UCEIS评分映射 (0-8分)
SEVERITY_MAPPING = {
    0: "缓解期",
    1: "轻度",
    2: "轻度",
    3: "中度",
    4: "中度",
    5: "重度",
    6: "重度",
    7: "重度",
    8: "重度"
}

# 响应模型
class PredictionResponse(BaseModel):
    task_id: str
    uceis_score: int
    severity: str
    confidence: float
    detailed_scores: Dict[str, Any]
    processing_time: float
    success: bool
    message: str

class HealthResponse(BaseModel):
    status: str
    model_loaded: bool
    device: str
    message: str

# 导入训练时使用的模型结构
sys.path.insert(0, r"D:\VSCodeProjects\UC_System\uc_model")
try:
    from models.fusion_model import EndoMultimodalModel
    logger.info("成功导入训练时使用的EndoMultimodalModel")
except ImportError as e:
    logger.error(f"无法导入训练模型: {e}")
    sys.exit(1)

def load_model() -> bool:
    """加载训练好的模型"""
    global model, device, tokenizer

    try:
        logger.info("开始加载UCEIS模型...")

        # 设备选择
        if torch.cuda.is_available():
            device = torch.device('cuda')
            logger.info(f"使用GPU: {torch.cuda.get_device_name(0)}")
        else:
            device = torch.device('cpu')
            logger.info("使用CPU进行推理")

        # 初始化tokenizer
        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        logger.info("BERT tokenizer加载成功")

        # 创建模型实例
        logger.info("初始化模型结构...")
        model = EndoMultimodalModel(
            image_pretrained=True,
            text_pretrained='bert-base-uncased',
            feature_dim=256,
            num_classes=8  # UCEIS评分0-7分（实际0-8分，但模型训练为8类）
        )
        logger.info("模型结构创建成功")

        # 加载训练好的权重
        model_path = Path(r"D:\VSCodeProjects\UC_System\uc_model\final_weights\uceis_best_model.pth")

        if not model_path.exists():
            logger.error(f"模型文件不存在: {model_path}")
            # 尝试其他可能的路径
            alt_paths = [
                Path(r"D:\VSCodeProjects\UC_System\uc_model\best_model_epoch_1.pth"),
                Path(r"D:\VSCodeProjects\UC_System\best_model_epoch_1.pth")
            ]
            for alt_path in alt_paths:
                if alt_path.exists():
                    model_path = alt_path
                    logger.info(f"找到替代模型文件: {model_path}")
                    break
            else:
                return False

        logger.info(f"加载模型权重: {model_path}")
        checkpoint = torch.load(model_path, map_location=device)

        # 加载模型状态字典
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'], strict=False)
            logger.info("成功加载model_state_dict")

            # 加载训练信息
            if 'epoch' in checkpoint:
                logger.info(f"训练轮数: {checkpoint['epoch']}")
            if 'best_f1' in checkpoint:
                logger.info(f"最佳F1分数: {checkpoint['best_f1']:.4f}")
        else:
            model.load_state_dict(checkpoint, strict=False)
            logger.info("直接加载checkpoint权重")

        # 移动到设备并设置为评估模式
        model.to(device)
        model.eval()

        # 计算参数量
        total_params = sum(p.numel() for p in model.parameters())
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)

        logger.info("模型加载成功！")
        logger.info(f"总参数量: {total_params:,}")
        logger.info(f"可训练参数: {trainable_params:,}")
        logger.info(f"模型设备: {device}")

        return True

    except Exception as e:
        logger.error(f"模型加载失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        return False

def preprocess_image(image_file) -> torch.Tensor:
    """图像预处理"""
    try:
        # 读取图像
        image = Image.open(image_file.file).convert('RGB')
        logger.info(f"成功加载图像: {image_file.filename}, 尺寸: {image.size}")

        # 检查图像质量
        if image.size[0] < 224 or image.size[1] < 224:
            logger.warning(f"图像尺寸过小: {image.size}")

        # 预处理变换 - 使用训练时相同的变换
        transform = transforms.Compose([
            transforms.Resize((224, 224)),  # 与训练时保持一致！
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                               std=[0.229, 0.224, 0.225])
        ])

        # 应用变换
        image_tensor = transform(image).unsqueeze(0)  # 添加batch维度
        logger.info(f"图像预处理完成，张量形状: {image_tensor.shape}")

        return image_tensor.to(device)

    except Exception as e:
        logger.error(f"图像预处理失败: {str(e)}")
        raise HTTPException(status_code=400, detail=f"图像预处理失败: {str(e)}")

def generate_image_description(image: Image.Image) -> str:
    """根据图像特征生成描述文本"""
    # 简单的图像分析
    width, height = image.size

    # 计算图像的基本统计
    img_array = np.array(image)
    avg_color = np.mean(img_array, axis=(0, 1))
    color_variance = np.var(img_array, axis=(0, 1))

    # 根据颜色特征判断
    if avg_color[0] > 200 and avg_color[1] < 100 and avg_color[2] < 100:
        # 偏红色 - 可能是出血
        return "endoscopic image showing red mucosa possibly indicating bleeding"
    elif avg_color[0] < 100 and avg_color[1] > 150 and avg_color[2] < 100:
        # 偏绿色 - 正常黏膜
        return "endoscopic image of normal green colon mucosa"
    elif np.max(color_variance) < 100:
        # 低方差 - 可能是溃疡或炎症
        return "endoscopic image showing ulceration or inflammation"
    else:
        # 默认描述
        return "endoscopic image of colon mucosa for UCEIS scoring assessment"

def predict_uceis_score(image_tensor: torch.Tensor, image_description: str) -> Dict[str, Any]:
    """预测UCEIS评分"""
    try:
        start_time = time.time()

        with torch.no_grad():
            # 使用动态生成的文本描述
            logger.info(f"使用描述: {image_description}")

            # 准备文本输入
            text_inputs = tokenizer(
                image_description,
                return_tensors="pt",
                padding=True,
                truncation=True,
                max_length=128
            )

            # 移动到正确设备
            for k, v in text_inputs.items():
                text_inputs[k] = v.to(device)

            # 模型推理
            logger.info("开始模型推理...")
            outputs = model(image_tensor, text_inputs)
            logger.info(f"模型输出形状: {outputs.shape if hasattr(outputs, 'shape') else type(outputs)}")

            # 处理输出
            if isinstance(outputs, (list, tuple)):
                logits = outputs[0] if len(outputs) > 0 else outputs
            else:
                logits = outputs

            # 确保是2D张量 [batch, classes]
            if logits.dim() == 1:
                logits = logits.unsqueeze(0)  # [classes] -> [1, classes]
            elif logits.dim() > 2:
                logits = logits.squeeze(0)  # 移除多余维度，但保留batch

            # 如果形状是 [batch, classes]，取第一个样本
            if logits.shape[0] > 1:
                logits = logits[0:1]  # 只取第一个样本
            elif logits.shape[0] == 0:
                raise ValueError("模型输出为空")

            # 获取预测结果 - 现在logits应该是 [1, 8]
            probabilities = F.softmax(logits, dim=1)  # 沿classes维度softmax
            predicted_class = torch.argmax(probabilities, dim=1).item()  # 获取预测类别

            # 转换为UCEIS评分（0-7分映射到UCEIS评分）
            uceis_score = min(predicted_class, 8)  # 确保不超过8分

            # 计算置信度
            confidence = probabilities[0, predicted_class].item()  # probabilities现在是 [1, 8]，取第0个样本的第predicted_class个元素

            # 获取所有概率
            all_probs = probabilities.cpu().numpy().tolist()

            logger.info(f"预测完成: UCEIS={uceis_score}, 置信度={confidence:.3f}")
            prob_str = [f'{i}:{p:.3f}' for i, p in enumerate(all_probs[0])]
            logger.info(f"概率分布: {prob_str}")

        processing_time = time.time() - start_time

        return {
            'uceis_score': uceis_score,
            'raw_prediction': predicted_class,
            'severity': SEVERITY_MAPPING.get(uceis_score, "未知"),
            'confidence': confidence,
            'all_probabilities': all_probs,
            'processing_time': processing_time
        }

    except Exception as e:
        logger.error(f"预测失败: {str(e)}")
        logger.error(f"错误详情: {traceback.format_exc()}")
        raise HTTPException(status_code=500, detail=f"预测失败: {str(e)}")

# 创建FastAPI应用
app = FastAPI(
    title="EndoSight-UC AI推理服务器 (修复版)",
    description="溃疡性结肠炎内镜严重程度指数(UCEIS)评分AI服务",
    version="2.1.0"
)

# 添加CORS中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.on_event("startup")
async def startup_event():
    """启动时加载模型"""
    success = load_model()
    if success:
        logger.info("✅ 推理服务器启动成功！")
    else:
        logger.error("❌ 模型加载失败，服务器将无法正常工作")

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    return HealthResponse(
        status="healthy" if model is not None else "unhealthy",
        model_loaded=model is not None,
        device=str(device) if device else "unknown",
        message="模型已加载" if model else "模型未加载"
    )

@app.post("/v1/predict", response_model=PredictionResponse)
async def predict_uceis(file: UploadFile = File(...)):
    """预测UCEIS评分"""
    if model is None:
        raise HTTPException(status_code=503, detail="模型未加载，请检查服务器状态")

    # 验证文件类型
    if not file.content_type or not file.content_type.startswith('image/'):
        raise HTTPException(status_code=400, detail="请上传图像文件")

    # 生成任务ID
    task_id = str(uuid.uuid4())

    try:
        # 预处理图像
        image_tensor = preprocess_image(file)

        # 读取原始图像用于分析
        file.file.seek(0)
        original_image = Image.open(file.file).convert('RGB')

        # 生成图像描述
        image_description = generate_image_description(original_image)

        # 预测
        result = predict_uceis_score(image_tensor, image_description)

        # 构建详细分数 - 基于实际推理结果
        detailed_scores = {
            'total': {
                'score': result['uceis_score'],
                'confidence': result['confidence'],
                'probabilities': result['all_probabilities']
            },
            'analysis': {
                'raw_prediction': result['raw_prediction'],
                'image_description': image_description,
                'model_confidence': result['confidence']
            }
        }

        # 构建响应
        response = PredictionResponse(
            task_id=task_id,
            uceis_score=result['uceis_score'],
            severity=result['severity'],
            confidence=result['confidence'],
            detailed_scores=detailed_scores,
            processing_time=result['processing_time'],
            success=True,
            message="预测成功"
        )

        logger.info(f"任务 {task_id} 预测完成: UCEIS={result['uceis_score']}, 置信度={result['confidence']:.3f}")
        return response

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"任务 {task_id} 处理失败: {str(e)}")
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")

@app.get("/")
async def root():
    """根路径"""
    return {
        "service": "EndoSight-UC AI推理服务器 (修复版)",
        "status": "running",
        "version": "2.1.0",
        "model_loaded": model is not None,
        "endpoints": {
            "/health": "健康检查",
            "/v1/predict": "UCEIS评分预测",
            "/docs": "API文档"
        }
    }

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=5001, log_level="info")